# -*- coding: utf-8 -*-
"""
系统测试脚本
用于验证检索方法对比评估系统是否正常工作
"""

import os
import sys
import json
from typing import List, Dict, Any
from loguru import logger

# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from config import OLLAMA_CONFIG
from data_manager import DataManager
from evaluator import Evaluator
from ollama_client import ollama_client
from retrieval_methods import (
    VectorRetriever, KeywordRetriever, HybridRetriever,
    KnowledgeGraphRetriever, RerankRetriever, MultiStepRetriever,
    StructuredRetriever
)


def setup_test_logging():
    """设置测试日志"""
    logger.remove()
    logger.add(
        sys.stdout,
        format="<green>{time:HH:mm:ss}</green> | <level>{level: <8}</level> | <level>{message}</level>",
        level="INFO"
    )


def create_test_data() -> List[Dict[str, Any]]:
    """创建测试数据"""
    return [
        {
            "id": 1,
            "question": "什么是机器学习？",
            "answer": "机器学习是人工智能的一个分支，通过算法让计算机从数据中学习，无需明确编程。",
            "category": 1,
            "difficulty": "medium",
            "type": "qa",
            "keywords": ["机器学习", "人工智能", "算法", "数据"]
        },
        {
            "id": 2,
            "question": "深度学习的原理是什么？",
            "answer": "深度学习使用多层神经网络来模拟人脑的学习过程，通过反向传播算法训练网络参数。",
            "category": 1,
            "difficulty": "hard",
            "type": "qa",
            "keywords": ["深度学习", "神经网络", "反向传播", "训练"]
        },
        {
            "id": 3,
            "question": "Python中如何定义函数？",
            "answer": "在Python中使用def关键字定义函数，语法为：def function_name(parameters): return value",
            "category": 2,
            "difficulty": "easy",
            "type": "code",
            "keywords": ["Python", "函数", "def", "参数"]
        },
        {
            "id": 4,
            "question": "什么是数据结构？",
            "answer": "数据结构是计算机存储、组织数据的方式，包括数组、链表、栈、队列、树、图等。",
            "category": 2,
            "difficulty": "medium",
            "type": "qa",
            "keywords": ["数据结构", "数组", "链表", "栈", "队列", "树", "图"]
        },
        {
            "id": 5,
            "question": "如何优化算法性能？",
            "answer": "算法优化可以从时间复杂度和空间复杂度两个方面考虑，使用更高效的数据结构和算法设计模式。",
            "category": 2,
            "difficulty": "hard",
            "type": "qa",
            "keywords": ["算法优化", "时间复杂度", "空间复杂度", "数据结构"]
        }
    ]


def test_ollama_connection() -> bool:
    """测试Ollama连接"""
    logger.info("🔗 测试Ollama连接...")
    
    try:
        if ollama_client.check_connection():
            logger.info("✅ Ollama连接正常")
            
            # 测试embedding功能
            test_text = "这是一个测试文本"
            embedding = ollama_client.get_embedding(test_text)
            if embedding and len(embedding) > 0:
                logger.info(f"✅ Embedding测试成功，维度: {len(embedding)}")
            else:
                logger.error("❌ Embedding测试失败")
            
            # 测试rerank功能
            test_docs = [
                {"content": "机器学习是人工智能的一个分支"},
                {"content": "深度学习使用神经网络"},
                {"content": "今天天气很好"}
            ]
            reranked = ollama_client.rerank_documents("什么是机器学习", test_docs)
            if reranked:
                logger.info(f"✅ Rerank测试成功，返回{len(reranked)}个结果")
                return True
            else:
                logger.error("❌ Rerank测试失败")
                return False
        else:
            logger.error("❌ Ollama连接失败")
            return False
            
    except Exception as e:
        logger.error(f"❌ Ollama测试异常: {e}")
        return False


def test_data_manager() -> bool:
    """测试数据管理器"""
    logger.info("📊 测试数据管理器...")
    
    try:
        data_manager = DataManager()
        test_data = create_test_data()
        
        # 测试数据处理功能
        categories = data_manager.get_questions_by_category(test_data, 1)
        if len(categories) == 2:  # 应该有2个分类1的问题
            logger.info("✅ 分类筛选功能正常")
        else:
            logger.error(f"❌ 分类筛选功能异常，期望2个，实际{len(categories)}个")
            return False
        
        # 测试关键词提取
        keywords = data_manager.extract_keywords("机器学习是人工智能的重要分支")
        if keywords and len(keywords) > 0:
            logger.info("✅ 关键词提取功能正常")
        else:
            logger.error("❌ 关键词提取功能异常")
            return False
        
        logger.info("✅ 数据管理器测试通过")
        return True
        
    except Exception as e:
        logger.error(f"❌ 数据管理器测试异常: {e}")
        return False


def test_retrievers() -> bool:
    """测试检索器"""
    logger.info("🔍 测试检索器...")
    
    test_data = create_test_data()
    test_question = "什么是机器学习？"
    
    # 测试各种检索器
    retrievers = {
        "向量检索": VectorRetriever,
        "关键词检索": KeywordRetriever,
        "混合检索": HybridRetriever,
        "知识图谱检索": KnowledgeGraphRetriever,
        "重排序检索": RerankRetriever,
        "多步推理检索": MultiStepRetriever,
        "结构化检索": StructuredRetriever
    }
    
    success_count = 0
    
    for name, retriever_class in retrievers.items():
        try:
            logger.info(f"  测试 {name}...")
            
            # 创建检索器实例
            retriever = retriever_class()
            
            # 初始化
            if retriever.initialize(test_data):
                # 执行检索
                result = retriever.retrieve(test_question, top_k=3)
                
                if result:
                    logger.info(f"  ✅ {name} 测试通过")
                    if hasattr(result, 'retrieved_docs') and result.retrieved_docs:
                        logger.info(f"     检索到 {len(result.retrieved_docs)} 个文档")
                    if hasattr(result, 'answer') and result.answer:
                        logger.info(f"     答案: {result.answer[:100]}...")
                    if hasattr(result, 'response_time'):
                        logger.info(f"     响应时间: {result.response_time:.3f}s")
                    success_count += 1
                else:
                    logger.warning(f"  ⚠️ {name} 返回空结果")
            else:
                logger.warning(f"  ⚠️ {name} 初始化失败")
                
        except Exception as e:
            logger.error(f"  ❌ {name} 测试异常: {e}")
    
    if success_count >= len(retrievers) // 2:  # 至少一半的检索器工作正常
        logger.info(f"✅ 检索器测试通过 ({success_count}/{len(retrievers)})")
        return True
    else:
        logger.error(f"❌ 检索器测试失败 ({success_count}/{len(retrievers)})")
        return False


def test_evaluator() -> bool:
    """测试评估器"""
    logger.info("📊 测试评估器...")
    
    try:
        test_data = create_test_data()
        
        # 创建评估器
        evaluator = Evaluator({
            'output_dir': 'test_results',
            'enable_detailed_analysis': False,
            'enable_visualization': False
        })
        
        # 初始化检索器（只测试几个快速的）
        if evaluator.initialize_retrievers(test_data[:3]):
            logger.info("✅ 评估器初始化成功")
            
            # 测试单个方法评估
            logger.info("  测试向量检索评估...")
            result = evaluator.evaluate_single_method("vector", test_data[:2])
            
            if result:
                logger.info("  ✅ 评估器测试成功")
                logger.info(f"     总查询数: {result.total_queries}")
                logger.info(f"     成功率: {result.metrics.get('success_rate', 0):.2%}")
                logger.info(f"     平均响应时间: {result.metrics.get('avg_response_time', 0):.3f}s")
                logger.info(f"     检索精度: {result.retrieval_precision:.3f}")
                logger.info(f"     检索召回: {result.retrieval_recall:.3f}")
                logger.info(f"     综合评分: {result.get_overall_score():.3f}")
                return True
            else:
                logger.error("❌ 评估器返回空结果")
                return False
        else:
            logger.error("❌ 评估器初始化失败")
            return False
            
    except Exception as e:
        logger.error(f"❌ 评估器测试异常: {e}")
        return False


def test_full_pipeline() -> bool:
    """测试完整流程"""
    logger.info("🚀 测试完整流程...")
    
    try:
        # 创建测试数据文件
        test_data = create_test_data()
        test_file = "test_data.json"
        
        with open(test_file, 'w', encoding='utf-8') as f:
            json.dump(test_data, f, ensure_ascii=False, indent=2)
        
        # 运行主程序（简化版）
        from main import run_evaluation
        
        summary = run_evaluation(
            methods=['keyword', 'vector'],  # 只测试两种快速方法
            sample_size=3,  # 小样本
            output_dir='test_results'
        )
        
        # 清理测试文件
        if os.path.exists(test_file):
            os.remove(test_file)
        
        if summary and summary.get('metrics_summary'):
            logger.info("✅ 完整流程测试通过")
            return True
        else:
            logger.error("❌ 完整流程测试失败")
            return False
            
    except Exception as e:
        logger.error(f"❌ 完整流程测试异常: {e}")
        return False


def run_system_test() -> bool:
    """运行系统测试"""
    logger.info("🧪 开始系统测试")
    logger.info("=" * 50)
    
    tests = [
        ("Ollama连接", test_ollama_connection),
        ("数据管理器", test_data_manager),
        ("检索器", test_retrievers),
        ("评估器", test_evaluator),
        ("完整流程", test_full_pipeline)
    ]
    
    passed_tests = 0
    total_tests = len(tests)
    
    for test_name, test_func in tests:
        logger.info(f"\n🔍 测试 {test_name}...")
        
        try:
            if test_func():
                passed_tests += 1
                logger.info(f"✅ {test_name} 测试通过")
            else:
                logger.error(f"❌ {test_name} 测试失败")
        except Exception as e:
            logger.error(f"❌ {test_name} 测试异常: {e}")
    
    logger.info("\n" + "=" * 50)
    logger.info(f"🎯 测试结果: {passed_tests}/{total_tests} 通过")
    
    if passed_tests == total_tests:
        logger.info("🎉 所有测试通过！系统运行正常")
        return True
    elif passed_tests >= total_tests * 0.8:  # 80%通过率
        logger.warning("⚠️ 大部分测试通过，系统基本正常")
        return True
    else:
        logger.error("❌ 多个测试失败，系统可能存在问题")
        return False


def main():
    """主函数"""
    setup_test_logging()
    
    logger.info("🧪 检索方法对比评估系统测试")
    logger.info(f"📍 工作目录: {os.getcwd()}")
    logger.info(f"🐍 Python版本: {sys.version}")
    
    # 检查基本环境
    logger.info("\n🔧 检查环境...")
    
    # 检查必要的目录
    dirs_to_check = ['data', 'results', 'logs']
    for dir_name in dirs_to_check:
        if not os.path.exists(dir_name):
            os.makedirs(dir_name)
            logger.info(f"📁 创建目录: {dir_name}")
    
    # 运行系统测试
    success = run_system_test()
    
    if success:
        logger.info("\n🎉 系统测试完成，可以开始正式评估！")
        logger.info("\n💡 使用以下命令开始评估:")
        logger.info("   python main.py --sample-size 10  # 小规模测试")
        logger.info("   python main.py                   # 完整评估")
        sys.exit(0)
    else:
        logger.error("\n❌ 系统测试失败，请检查配置和环境")
        logger.info("\n🔧 故障排除建议:")
        logger.info("   1. 确保Ollama服务正在运行")
        logger.info("   2. 检查qwq:latest模型是否已下载")
        logger.info("   3. 安装所有依赖包: pip install -r requirements.txt")
        logger.info("   4. 检查网络连接")
        sys.exit(1)


if __name__ == "__main__":
    main()